﻿//#define NATIVE_WEBSOCKET

using System;
using System.Collections.Generic;
using System.Threading;
using Cysharp.Threading.Tasks;

#if !NATIVE_WEBSOCKET && (UNITY_EDITOR || !UNITY_WEBGL)
using System.Buffers;
using System.IO;
using System.Net.WebSockets;
#endif

namespace KuiHuaBaoDian.Services.Net {

    /// <summary>
    /// 网络连接
    /// </summary>
    public interface INetConnection {

        public delegate void StatusChangedDelegate(Status status);

        public enum Status {
            Success,
            Cancel,
            Close,
            Error,
        }

        event StatusChangedDelegate OnStatusChanged;

        NetConnectionHandle Handle { get; }

        bool IsConnected { get; }

        UniTask Close(bool needFire);

        UniTask Connect(string url, CancellationToken cancellationToken);

        byte[] GetReceived();

        UniTask Send(byte[] bytes);

        void Dispose();
    }

#if !NATIVE_WEBSOCKET && (UNITY_EDITOR || !UNITY_WEBGL)
    [Implement(typeof(INetConnection))]
    internal sealed class KHBD_WebSocketNetConnection : INetConnection {

        public event INetConnection.StatusChangedDelegate OnStatusChanged;

        private ClientWebSocket m_Socket;
        private Queue<byte[]> m_BytesQueue = new();
        private object m_BytesQueueLocker = new();

        public NetConnectionHandle Handle { get; private set; } = new NetConnectionHandle { GUID = Guid.NewGuid() };

        public bool IsConnected {
            get {
                if (m_Socket != null) {
                    return !m_IsClosing && m_Socket.State == WebSocketState.Open;
                } else {
                    return false;
                }
            }
        }

        public async UniTask Connect(string url, CancellationToken cancellationToken) {
            UnityEngine.Assertions.Assert.IsNull(m_Socket);

            if (KHBD.Context.IsCancellationRequested(cancellationToken)) {
                OnStatusChanged?.Invoke(INetConnection.Status.Cancel);
                return;
            }

            try {
                m_Socket = new ClientWebSocket();
                var uri = new Uri(url);
                await m_Socket.ConnectAsync(uri, cancellationToken);
                OnStatusChanged?.Invoke(INetConnection.Status.Success);
            } catch (OperationCanceledException) {
                m_Socket.Dispose();
                m_Socket = null;
                OnStatusChanged?.Invoke(INetConnection.Status.Cancel);
            } catch (Exception e) {
                m_Socket.Dispose();
                m_Socket = null;
                UnityEngine.Debug.LogWarning(e.Message);
                OnStatusChanged?.Invoke(INetConnection.Status.Error);
            }
            if (m_Socket != null) {
                BeginReceive().Forget();
            }
        }

        private async UniTaskVoid BeginReceive() {
            while (IsConnected && KHBD.Context.IsApplicationPlaying2) {
                try {
                    var result = default(WebSocketReceiveResult);
                    using var memoryStream = new MemoryStream();

                    while (IsConnected && (result == null || !result.EndOfMessage) && KHBD.Context.IsApplicationPlaying) {
                        var shared = ArrayPool<byte>.Shared;
                        var rented = shared.Rent(8192);
                        var buffer = new ArraySegment<byte>(rented);
                        try {
                            result = await m_Socket.ReceiveAsync(buffer, CancellationToken.None);
                            memoryStream.Write(buffer.Array, buffer.Offset, result.Count);
                        } catch {
                            result = new WebSocketReceiveResult(0, WebSocketMessageType.Close, true);
                        }
                        shared.Return(rented);
                        // TODO: 连续字节
                    }

                    if (result.MessageType == WebSocketMessageType.Text || result.MessageType == WebSocketMessageType.Binary) {
                        memoryStream.Seek(0, SeekOrigin.Begin);
                        lock (m_BytesQueueLocker) {
                            m_BytesQueue.Enqueue(memoryStream.ToArray());
                        }
                    } else if (result.MessageType == WebSocketMessageType.Close) {
                        if (!m_IsClosing) {
                            OnStatusChanged?.Invoke(INetConnection.Status.Error);
                        }
                        break;
                    }
                } catch (Exception e) {
                    if (!m_IsClosing) {
                        await Close(false);
                        UnityEngine.Debug.LogException(e);
                        OnStatusChanged?.Invoke(INetConnection.Status.Error);
                    }
                }
            }
            await UniTask.Yield();
        }

        public byte[] GetReceived() {
            if (m_BytesQueue.Count > 0) {
                lock (m_BytesQueueLocker) {
                    return m_BytesQueue.Dequeue();
                }
            }
            return null;
        }

        public async UniTask Send(byte[] bytes) {
            UnityEngine.Assertions.Assert.IsNotNull(bytes);

            if (IsConnected) {
                try {
                    await m_Socket.SendAsync(new ArraySegment<byte>(bytes), WebSocketMessageType.Binary, true, CancellationToken.None);
                } catch (Exception e) {
                    UnityEngine.Debug.LogException(e);
                }
            }
        }

        private bool m_IsClosing = false;
        public async UniTask Close(bool needFire) {
            if (m_Socket == null || m_IsClosing) {
                return;
            }

            m_IsClosing = true;
            try {
                await m_Socket.CloseAsync(WebSocketCloseStatus.NormalClosure, string.Empty, CancellationToken.None);
            } catch {
            }
            m_IsClosing = false;

            m_Socket.Abort();
            m_Socket.Dispose();
            m_Socket = null;

            lock (m_BytesQueueLocker) {
                m_BytesQueue.Clear();
            }

            if (needFire) {
                OnStatusChanged?.Invoke(INetConnection.Status.Close);
            }
        }

        public void Dispose() {
            OnStatusChanged = null;

            lock (m_BytesQueueLocker) {
                m_BytesQueue = null;
            }

            m_BytesQueueLocker = null;

            Handle = NetConnectionHandle.None;
        }
    }
#endif

#if NATIVE_WEBSOCKET || !UNITY_EDITOR && UNITY_WEBGL
    [Implement(typeof(INetConnection))]
    internal sealed class KHBD_NativeWebSocketNetConnection : INetConnection {

        public NetConnectionHandle Handle { get; private set; } = new NetConnectionHandle { GUID = Guid.NewGuid() };

        public bool IsConnected => m_Socket.State == NativeWebSocket.WebSocketState.Open;

        public event INetConnection.StatusChangedDelegate OnStatusChanged;

        private NativeWebSocket.WebSocket m_Socket;
        private Queue<byte[]> m_BytesQueue = new();

        private bool m_IsClosing = false;
        public async UniTask Close(bool needFire) {
            if (m_Socket == null || m_IsClosing) {
                return;
            }
            m_IsClosing = true;
            try {
                m_Socket.OnClose -= OnClose;
                await m_Socket.Close();
            } catch {
            }
            m_IsClosing = false;
            Closed(needFire);
        }

        private void Closed(bool needFire) {
            m_Socket = null;
            m_BytesQueue.Clear();
            if (needFire) {
                OnStatusChanged?.Invoke(INetConnection.Status.Close);
            }
        }

        public async UniTask Connect(string url, CancellationToken cancellationToken) {
            UnityEngine.Assertions.Assert.IsNull(m_Socket);

            if (cancellationToken.IsCancellationRequested) {
                OnStatusChanged?.Invoke(INetConnection.Status.Cancel);
                return;
            }

            m_Socket = new NativeWebSocket.WebSocket(url);
            var isOpened = false;
            var isFailed = false;
            m_Socket.OnOpen += () => {
                isOpened = true;
            };
            m_Socket.OnError += message => {
                if (isOpened) {
                    isOpened = false;
                    isFailed = true;
                    UnityEngine.Debug.LogWarning(message);
                    Close(false).Forget();
                    OnStatusChanged?.Invoke(INetConnection.Status.Error);
                } else {
                    isFailed = true;
                }
            };
            m_Socket.OnClose += code => {
                UnityEngine.Debug.Log(code);
            };
            m_Socket.OnMessage += (bytes) => {
                m_BytesQueue.Enqueue(bytes);
            };
            await m_Socket.Connect();
            if (cancellationToken.IsCancellationRequested) {
                m_Socket.CancelConnection();
                m_Socket = null;
                OnStatusChanged?.Invoke(INetConnection.Status.Cancel);
                return;
            }
            while (!isOpened && !isFailed && KHBD.Context.IsApplicationPlaying) {
                await UniTask.Yield();
            }
            if (cancellationToken.IsCancellationRequested) {
                if (isOpened) {
                    m_Socket.CancelConnection();
                }
                m_Socket = null;
                OnStatusChanged?.Invoke(INetConnection.Status.Cancel);
                return;
            }
            if (isOpened) {
                m_Socket.OnClose += OnClose;
                OnStatusChanged?.Invoke(INetConnection.Status.Success);
            }
        }

        private void OnClose(NativeWebSocket.WebSocketCloseCode _) {
            Closed(false);
            OnStatusChanged?.Invoke(INetConnection.Status.Error);
        }

        public void Dispose() {
            m_Socket = null;
            m_BytesQueue = null;
            OnStatusChanged = null;
            Handle = NetConnectionHandle.None;
        }

        public byte[] GetReceived() {
            if (m_BytesQueue.Count > 0) {
                return m_BytesQueue.Dequeue();
            }
            return null;
        }

        public async UniTask Send(byte[] bytes) {
            await m_Socket.Send(bytes);
        }
    }
#endif
}
